# coding=utf-8

'''
9-18作业：
在室温大气压下，把氢气看做氢气分子的理想气体，服从麦克斯韦速率分布函数：
$$
f(v) = 4 \pi v^2 ( \frac{m}{2\pi kT})^{3/2} exp( -\frac{mv^2}{2kT})\\
= \frac{4}{\sqrt{\pi}} (\frac{v^2}{v^3_p}) exp(-\frac{v^2}{v_p^2})
$$
其中最概然速率是$v_p = \sqrt{\frac{2kT}{m}}$（室温大气压下氢气分子最概然速率为$v_p = 1578m/s$）。

如果以$v_p$为速率单位，即设$v_p=1$，绘制$f(v)\sim v$的曲线。

试求：速率在以下区间内的分子数占总分子数的比例
- $0 \sim v_p$，
- $0 \sim 3.3v_p$,
要求：误差不超过$0.01\%$.
'''

import numpy as np
import matplotlib.pyplot as plt

# 被积函数
def f(v):
   # return v*v;
    return 4 / np.sqrt(np.pi) * v * v * np.exp( - v*v)

# trapzoid integration of f(x) on [a,b], with n divisions
def trapzoid1(a, b, f, n):
    s = 0
    h = (b-a)/n # 步长
    for i in range(n): #range(n) = [0,1,...,n-1]
        s += ( f(a + i*h) + f(a + (i+1)*h) ) * h / 2
    return s

# trapzoid integration of f(x) on [a,b], with error <= precision
def trapzoid2(a, b, f, precision): # 返还一个元组 ( 积分值，误差 )
    n = 10
    error1 = 1000 * precision
    NIn = 1 # global variables
    N2n = 1
    while( error1 > precision ):
        NIn = trapzoid1(a, b, f, n)
        NI2n = trapzoid1(a, b, f, 2*n)
        error1 = (NIn - NI2n)/3
        #print(" n = ", n, " integral = ", Delta2n)
        #print(" error1 = ", error1, "precision = ", precision)
        n = 2*n
    return NI2n, (NIn - NI2n)/3

# (1) 作图： f(v) ~ v
xdata = np.arange(0, 5, 0.1)
ydata = f( xdata )
plt.plot(xdata, ydata, label = "Ideal Hydrogen gas: maxwell distribution")
plt.xlabel("v(unit: $v_p = 1578m/s$)",fontsize=15)
plt.ylabel("f(v)",fontsize=15)
plt.legend()
plt.savefig("MaxwellDistribution.png")
plt.show()

# (2) [0,v_p] integration
x = trapzoid2(0,1,f,0.00001)
print( "Probability of v in [ 0, v_p ] = ", x[0], ", error ~ ", x[1] )

# (3) [0,3.3v_p] integration
x = trapzoid2(0,3.3,f,0.00001)
print( "Probability of v in [ 0, 3.3 v_p ] = ", x[0], ", error ~ ", x[1])
